from setting import backtest
from setting import data
import polars as pl
import pickle
from multiprocessing import Pool
from setting.crate import *
# 策略名称 - 用于生成测试报告
strategy_name = "Alice_baby"
# 参数组合
parameter_combinations = [(5, 10), (4, 8), (10, 20)]

# 合约池
base_data_symbol = ["CFFEX_IF00_300"]  # 基准合约列表(用于时间轴)
trade_data_symbol = ["CFFEX_IF00_300","CFFEX_IC00_300","CFFEX_IH00_300",] 
extra_data_symbol = [] 

# 加载数据
base_data, all_trade_data = data.get_data(
    base_data_symbol, 
    trade_data_symbol + extra_data_symbol, 
    "20250101", 
    "20251031"
)

# 分割数据
trade_data = {k: v for k, v in all_trade_data.items() if k in trade_data_symbol}
extra_data = {k: v for k, v in all_trade_data.items() if k in extra_data_symbol}

# 回测参数配置
initial_cash = 1000 * 10000.0  # 初始资金:1000万
fee_rate = 3.0 / 10000  # 手续费率:万分之三

# 加载合约参数
with open("setting/volume_multiple.pkl", "rb") as f:
    loaded_volume = pickle.load(f)
with open("setting/margin_ratio.pkl", "rb") as f:
    loaded_margin = pickle.load(f)

deposit_rates = loaded_margin
multipliers = loaded_volume

# 一次性添加所有字段
for symbol, data in trade_data.items():
    lazy_plan = data.lazy()
    # 一次性添加所有均线字段
    for i in [5, 10, 4, 8, 10, 20]:
        lazy_plan = lazy_plan.with_columns([
            pl.col("close").rolling_mean(i).alias(f"ma_{i}"),
            pl.col("close").rolling_mean(i).shift(1).alias(f"prev_ma_{i}")
        ])
    # 最后一次性执行
    trade_data[symbol] = lazy_plan.collect()
class Strategy(PythonStrategy):
    """
    多参数优化的均线交叉策略
    """

    def __init__(self, backtest_instance, trade_data, extra_data, n1, n2):
        """
        初始化策略
        Args:
            backtest_instance: 回测引擎实例
            trade_data: 交易数据
            extra_data: 额外数据
            n1: 参数1
            n2: 参数2
        """
        super().__init__(backtest_instance, trade_data, extra_data)
        self.n1 = n1
        self.n2 = n2

    def on_bar(self, timestamp):
        """K线回调函数 - 每个时间戳调用一次"""
        for symbol in self.trade_data.keys():
            recent_data = self.get_recent_data(symbol, timestamp, lookback=max(self.n1,self.n2)+1, require_exact_match=True)
            if recent_data is not None:
                self.calculate_ma_cross_signal(extract_symbol_base(symbol), recent_data, timestamp)

    def calculate_ma_cross_signal(self, symbol, recent_data, timestamp):
        """
        均线交叉信号计算函数
        
        Args:
            symbol: 品种代码
            recent_data: 最近的数据切片
            timestamp: 时间戳
        """
        # 直接从预计算的数据中读取指标值
        current_fast = recent_data["ma_"+str(self.n1)][-1]  # 当前fast均线
        current_slow = recent_data["ma_"+str(self.n2)][-1]  # 当前slow均线
        prev_fast = recent_data["prev_ma_"+str(self.n1)][-1]  # 前一个fast周期均线
        prev_slow = recent_data["prev_ma_"+str(self.n2)][-1]  # 前一个slow周期均线

        # 获取当前持仓
        current_holding = self.backtest_instance.long_position.get(symbol, 0)

        # 金叉信号开多
        if (prev_fast <= prev_slow and current_fast > current_slow 
            and current_holding == 0):
            self.backtest_instance.buy(symbol, 1, recent_data["close"][-1])

        # 死叉信号平仓
        elif (prev_fast >= prev_slow and current_fast < current_slow 
              and current_holding > 0):
            self.backtest_instance.sell(symbol, current_holding, recent_data["close"][-1])


def run_single_strategy(params):
    """单个策略回测任务"""
    params_name = str(params)
    
    # 初始化回测引擎
    my_backtest = backtest.BackTest(initial_cash, fee_rate, deposit_rates, multipliers, trade_data)
    
    # 创建策略实例 - 使用子类Strategy
    strategy = Strategy(my_backtest, trade_data, extra_data, params[0], params[1])
    
    # 准备回测时间序列
    datetime_strings = base_data.get_column("datetime").dt.strftime("%Y-%m-%d %H:%M:%S")
    
    # 启动回测
    results = backtest.batch_backtest(
        strategy,
        my_backtest,
        datetime_strings,
        base_data.get_column("timestamps"),
        strategy_name,
        params_name,
    )
    
    return {"strategy_name": strategy_name,"params_name": params_name,"results": results}


def main():
    """主函数 - 多进程参数优化"""

    print(f"🚀 开始多进程参数优化，共 {len(parameter_combinations)} 个参数组合")
    
    # 多进程执行
    with Pool(4) as pool:
        results = pool.map(run_single_strategy, parameter_combinations)
    
    # 导出CSV
    export_results_to_csv(results)


if __name__ == "__main__":
    main()